#define MSAES_C
#include <stdio.h>
#include <string.h>
#include <libmslog/mslog.h>
#include "mscommon.h"   
#include "msaes.h"   

#define FLAG "msaes"

typedef struct _ST_CA_AES_CONTEXT
{
	ms_s32 nr;         	/*!<  number of rounds  */
	ms_pu32 rk;      	/*!<  AES round keys    */
	ms_u32 buf[68];    	/*!<  unaligned data    */
}ST_CA_AES_CONTEXT;
#define PAS_RX_LEN 256

//加密解密 
/*
 * 32-bit INT32eger manipulation macros (little endian)
 */
#define GET_ULONG_LE(n,b,i)                             \
{                                                       \
    (n) = ( (ms_u32) (b)[(i)    ]       )        \
        | ( (ms_u32) (b)[(i) + 1] <<  8 )        \
        | ( (ms_u32) (b)[(i) + 2] << 16 )        \
        | ( (ms_u32) (b)[(i) + 3] << 24 );       \
}

#define PUT_ULONG_LE(n,b,i)                             \
{                                                       \
    (b)[(i)    ] = (ms_u08) ( (n)       );       \
    (b)[(i) + 1] = (ms_u08) ( (n) >>  8 );       \
    (b)[(i) + 2] = (ms_u08) ( (n) >> 16 );       \
    (b)[(i) + 3] = (ms_u08) ( (n) >> 24 );       \
}
/*
 * Tables generation code
 */
#define ROTL8(x) ( ( x << 8 ) & 0xFFFFFFFF ) | ( x >> 24 )
#define XTIME(x) ( ( x << 1 ) ^ ( ( x & 0x80 ) ? 0x1B : 0x00 ) )
#define MUL(x,y) ( ( x && y ) ? pow[(log[x]+log[y]) % 255] : 0 )

#define AES_FROUND(X0,X1,X2,X3,Y0,Y1,Y2,Y3)     \
{                                               \
	X0 = *RK++ ^ FT0[ ( Y0       ) & 0xFF ] ^   \
		FT1[ ( Y1 >>  8 ) & 0xFF ] ^   \
		FT2[ ( Y2 >> 16 ) & 0xFF ] ^   \
		FT3[ ( Y3 >> 24 ) & 0xFF ];    \
	                        \
	X1 = *RK++ ^ FT0[ ( Y1       ) & 0xFF ] ^   \
		FT1[ ( Y2 >>  8 ) & 0xFF ] ^   \
		FT2[ ( Y3 >> 16 ) & 0xFF ] ^   \
		FT3[ ( Y0 >> 24 ) & 0xFF ];    \
	                        \
	X2 = *RK++ ^ FT0[ ( Y2       ) & 0xFF ] ^   \
		FT1[ ( Y3 >>  8 ) & 0xFF ] ^   \
		FT2[ ( Y0 >> 16 ) & 0xFF ] ^   \
		FT3[ ( Y1 >> 24 ) & 0xFF ];    \
	                        \
	X3 = *RK++ ^ FT0[ ( Y3       ) & 0xFF ] ^   \
		FT1[ ( Y0 >>  8 ) & 0xFF ] ^   \
		FT2[ ( Y1 >> 16 ) & 0xFF ] ^   \
		FT3[ ( Y2 >> 24 ) & 0xFF ];    \
}

#define AES_RROUND(X0,X1,X2,X3,Y0,Y1,Y2,Y3)     \
{                                               \
	X0 = *RK++ ^ RT0[ ( Y0       ) & 0xFF ] ^   \
		RT1[ ( Y3 >>  8 ) & 0xFF ] ^   \
		RT2[ ( Y2 >> 16 ) & 0xFF ] ^   \
		RT3[ ( Y1 >> 24 ) & 0xFF ];    \
	                        \
	X1 = *RK++ ^ RT0[ ( Y1       ) & 0xFF ] ^   \
		RT1[ ( Y0 >>  8 ) & 0xFF ] ^   \
		RT2[ ( Y3 >> 16 ) & 0xFF ] ^   \
		RT3[ ( Y2 >> 24 ) & 0xFF ];    \
	                        \
	X2 = *RK++ ^ RT0[ ( Y2       ) & 0xFF ] ^   \
		RT1[ ( Y1 >>  8 ) & 0xFF ] ^   \
		RT2[ ( Y0 >> 16 ) & 0xFF ] ^   \
		RT3[ ( Y3 >> 24 ) & 0xFF ];    \
	                        \
	X3 = *RK++ ^ RT0[ ( Y3       ) & 0xFF ] ^   \
		RT1[ ( Y2 >>  8 ) & 0xFF ] ^   \
		RT2[ ( Y1 >> 16 ) & 0xFF ] ^   \
		RT3[ ( Y0 >> 24 ) & 0xFF ];    \
}

/*
 * Forward S-box & tables
 */
static ms_u08 FSb[256];
static ms_u32 FT0[256];
static ms_u32 FT1[256];
static ms_u32 FT2[256];
static ms_u32 FT3[256];

/*
 * Reverse S-box & tables
 */
static ms_u08 RSb[256];
static ms_u32 RT0[256];
static ms_u32 RT1[256];
static ms_u32 RT2[256];
static ms_u32 RT3[256];

/*
 * Round constants
 */
static ms_u32 RCON[10];

static void msaes_innerapi_gentable( void )
{
	ms_s32 i, x, y, z;
	ms_s32 pow[256];
	ms_s32 log[256];

	/*
	* compute pow and log tables over GF(2^8)
	*/
	for( i = 0, x = 1; i < 256; i++ ){
		pow[i] = x;
		log[x] = i;
		x = ( x ^ XTIME( x ) ) & 0xFF;
	}

	/*
	* calculate the round constants
	*/
	for( i = 0, x = 1; i < 10; i++ ){
		RCON[i] = (ms_u32) x;
		x = XTIME( x ) & 0xFF;
	}

	/*
	* generate the forward and reverse S-boxes
	*/
	FSb[0x00] = 0x63;
	RSb[0x63] = 0x00;

	for( i = 1; i < 256; i++ ){
		x = pow[255 - log[i]];
		y  = x; y = ( (y << 1) | (y >> 7) ) & 0xFF;
		x ^= y; y = ( (y << 1) | (y >> 7) ) & 0xFF;
		x ^= y; y = ( (y << 1) | (y >> 7) ) & 0xFF;
		x ^= y; y = ( (y << 1) | (y >> 7) ) & 0xFF;
		x ^= y ^ 0x63;

		FSb[i] = (ms_u08) x;
		RSb[x] = (ms_u08) i;
	}

	/*
	* generate the forward and reverse tables
	*/
	for( i = 0; i < 256; i++ ){
		x = FSb[i];
		y = XTIME( x ) & 0xFF;
		z =  ( y ^ x ) & 0xFF;

		FT0[i] = ( (ms_u32) y  ) ^
				( (ms_u32) x << 8 ) ^
				( (ms_u32) x << 16 ) ^
				( (ms_u32) z << 24 );

		FT1[i] = ROTL8( FT0[i] );
		FT2[i] = ROTL8( FT1[i] );
		FT3[i] = ROTL8( FT2[i] );

		x = RSb[i];

		RT0[i] = ( (ms_u32) MUL( 0x0E, x )       ) ^
				( (ms_u32) MUL( 0x09, x ) <<  8 ) ^
				( (ms_u32) MUL( 0x0D, x ) << 16 ) ^
				( (ms_u32) MUL( 0x0B, x ) << 24 );

		RT1[i] = ROTL8( RT0[i] );
		RT2[i] = ROTL8( RT1[i] );
		RT3[i] = ROTL8( RT2[i] );
	}
}
 
/*
 * AES key schedule (encryption)
 */
static void msaes_innerapi_setEncKey( ST_CA_AES_CONTEXT *ctx, ms_u08 *key, ms_s32 keysize )
{
	ms_s32 i;
	ms_u32 *RK;

	msaes_innerapi_gentable();
	switch( keysize ){
		case 128: ctx->nr = 10; break;
		case 192: ctx->nr = 12; break;
		case 256: ctx->nr = 14; break;
		default : return;
	}

	ctx->rk = RK = ctx->buf;
	if(RK==ms_null){
		printf("%s %d ms_null\n",__FUNCTION__,__LINE__);
		return;
	}

	for( i = 0; i < (keysize >> 5); i++ ){
	    if((i>68)||((i<<2)>64)){
           printf("%s %d fail\n",__FUNCTION__,__LINE__);
           return;
        }
		GET_ULONG_LE( RK[i], key, i << 2 );
	}

	switch( ctx->nr ){
		case 10:
			for( i = 0; i < 10; i++, RK += 4 ){
				RK[4]  = RK[0] ^ RCON[i] ^
						( (ms_u32) FSb[ ( RK[3] >>  8 ) & 0xFF ]       ) ^
						( (ms_u32) FSb[ ( RK[3] >> 16 ) & 0xFF ] <<  8 ) ^
						( (ms_u32) FSb[ ( RK[3] >> 24 ) & 0xFF ] << 16 ) ^
						( (ms_u32) FSb[ ( RK[3]       ) & 0xFF ] << 24 );

				RK[5]  = RK[1] ^ RK[4];
				RK[6]  = RK[2] ^ RK[5];
				RK[7]  = RK[3] ^ RK[6];
			}
			break;
		case 12:
			for( i = 0; i < 8; i++, RK += 6 ){
				RK[6]  = RK[0] ^ RCON[i] ^
						( (ms_u32) FSb[ ( RK[5] >>  8 ) & 0xFF ]   ) ^
						( (ms_u32) FSb[ ( RK[5] >> 16 ) & 0xFF ] <<  8 ) ^
						( (ms_u32) FSb[ ( RK[5] >> 24 ) & 0xFF ] << 16 ) ^
						( (ms_u32) FSb[ ( RK[5]       ) & 0xFF ] << 24 );

				RK[7]  = RK[1] ^ RK[6];
				RK[8]  = RK[2] ^ RK[7];
				RK[9]  = RK[3] ^ RK[8];
				RK[10] = RK[4] ^ RK[9];
				RK[11] = RK[5] ^ RK[10];
			}
			break;
		case 14:
			for( i = 0; i < 7; i++, RK += 8 ){
				RK[8]  = RK[0] ^ RCON[i] ^
						( (ms_u32) FSb[ ( RK[7] >>  8 ) & 0xFF ]       ) ^
						( (ms_u32) FSb[ ( RK[7] >> 16 ) & 0xFF ] <<  8 ) ^
						( (ms_u32) FSb[ ( RK[7] >> 24 ) & 0xFF ] << 16 ) ^
						( (ms_u32) FSb[ ( RK[7]       ) & 0xFF ] << 24 );

				RK[9]  = RK[1] ^ RK[8];
				RK[10] = RK[2] ^ RK[9];
				RK[11] = RK[3] ^ RK[10];

				RK[12] = RK[4] ^
						( (ms_u32) FSb[ ( RK[11]       ) & 0xFF ]       ) ^
						( (ms_u32) FSb[ ( RK[11] >>  8 ) & 0xFF ] <<  8 ) ^
						( (ms_u32) FSb[ ( RK[11] >> 16 ) & 0xFF ] << 16 ) ^
						( (ms_u32) FSb[ ( RK[11] >> 24 ) & 0xFF ] << 24 );

				RK[13] = RK[5] ^ RK[12];
				RK[14] = RK[6] ^ RK[13];
				RK[15] = RK[7] ^ RK[14];
			}
			break;
		default:
			break;
	}
}

/*
 * AES key schedule (decryption)
 */
static void msaes_innerapi_setDecKey( ST_CA_AES_CONTEXT *ctx, ms_u08 *key, ms_s32 keysize )
{
	ms_s32 i, j;
	ST_CA_AES_CONTEXT cty;
	ms_u32 *RK;
	ms_u32 *SK;

	switch( keysize ){
		case 128: ctx->nr = 10; break;
		case 192: ctx->nr = 12; break;
		case 256: ctx->nr = 14; break;
		default : return;
	}

	ctx->rk = RK = ctx->buf;
    if(RK==ms_null){
       printf("%s %d ms_null\n",__FUNCTION__,__LINE__);
       return;
    }

	msaes_innerapi_setEncKey( &cty, key, keysize );
	SK = cty.rk + cty.nr * 4;

	*RK++ = *SK++;
	*RK++ = *SK++;
	*RK++ = *SK++;
	*RK++ = *SK++;

	for( i = ctx->nr - 1, SK -= 8; i > 0; i--, SK -= 8 ){
		for( j = 0; j < 4; j++, SK++ ){
			*RK++ = RT0[ FSb[ ( *SK       ) & 0xFF ] ] ^
			RT1[ FSb[ ( *SK >>  8 ) & 0xFF ] ] ^
			RT2[ FSb[ ( *SK >> 16 ) & 0xFF ] ] ^
			RT3[ FSb[ ( *SK >> 24 ) & 0xFF ] ];
		}
	}

	*RK++ = *SK++;
	*RK++ = *SK++;
	*RK++ = *SK++;
	*RK++ = *SK++;

	memset( &cty, 0, sizeof( ST_CA_AES_CONTEXT ) );
}

/*
 * AES-ECB block encryption/decryption
 */
static void msaes_innerapi_cryptECB( ST_CA_AES_CONTEXT *ctx, ms_s32 mode,
                    							ms_u08 input[PAS_RX_LEN], ms_u08 output[PAS_RX_LEN] )
{
	ms_s32 i;
	ms_u32 *RK, X0, X1, X2, X3, Y0, Y1, Y2, Y3;

	RK = ctx->rk;
    if(RK==ms_null){
       printf("%s %d ms_null\n",__FUNCTION__,__LINE__);
       return;
    }

	GET_ULONG_LE( X0, input,  0 ); X0 ^= *RK++;
	GET_ULONG_LE( X1, input,  4 ); X1 ^= *RK++;
	GET_ULONG_LE( X2, input,  8 ); X2 ^= *RK++;
	GET_ULONG_LE( X3, input, 12 ); X3 ^= *RK++;

	if( mode == CRYPT_TYPE_DECRYPT ){
		for( i = (ctx->nr >> 1) - 1; i > 0; i-- ){
			AES_RROUND( Y0, Y1, Y2, Y3, X0, X1, X2, X3 );
			AES_RROUND( X0, X1, X2, X3, Y0, Y1, Y2, Y3 );
		}

		AES_RROUND( Y0, Y1, Y2, Y3, X0, X1, X2, X3 );

		X0 = *RK++ ^ \
		( (ms_u32) RSb[ ( Y0       ) & 0xFF ]       ) ^
		( (ms_u32) RSb[ ( Y3 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) RSb[ ( Y2 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) RSb[ ( Y1 >> 24 ) & 0xFF ] << 24 );

		X1 = *RK++ ^ \
		( (ms_u32) RSb[ ( Y1       ) & 0xFF ]       ) ^
		( (ms_u32) RSb[ ( Y0 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) RSb[ ( Y3 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) RSb[ ( Y2 >> 24 ) & 0xFF ] << 24 );

		X2 = *RK++ ^ \
		( (ms_u32) RSb[ ( Y2       ) & 0xFF ]       ) ^
		( (ms_u32) RSb[ ( Y1 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) RSb[ ( Y0 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) RSb[ ( Y3 >> 24 ) & 0xFF ] << 24 );

		X3 = *RK++ ^ \
		( (ms_u32) RSb[ ( Y3       ) & 0xFF ]       ) ^
		( (ms_u32) RSb[ ( Y2 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) RSb[ ( Y1 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) RSb[ ( Y0 >> 24 ) & 0xFF ] << 24 );
	}else {	/* AS_CRYPT_TYPE_DECRYPT */
		for( i = (ctx->nr >> 1) - 1; i > 0; i-- ){
			AES_FROUND( Y0, Y1, Y2, Y3, X0, X1, X2, X3 );
			AES_FROUND( X0, X1, X2, X3, Y0, Y1, Y2, Y3 );
		}

		AES_FROUND( Y0, Y1, Y2, Y3, X0, X1, X2, X3 );

		X0 = *RK++ ^ \
		( (ms_u32) FSb[ ( Y0       ) & 0xFF ]       ) ^
		( (ms_u32) FSb[ ( Y1 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) FSb[ ( Y2 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) FSb[ ( Y3 >> 24 ) & 0xFF ] << 24 );

		X1 = *RK++ ^ \
		( (ms_u32) FSb[ ( Y1       ) & 0xFF ]       ) ^
		( (ms_u32) FSb[ ( Y2 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) FSb[ ( Y3 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) FSb[ ( Y0 >> 24 ) & 0xFF ] << 24 );

		X2 = *RK++ ^ \
		( (ms_u32) FSb[ ( Y2       ) & 0xFF ]       ) ^
		( (ms_u32) FSb[ ( Y3 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) FSb[ ( Y0 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) FSb[ ( Y1 >> 24 ) & 0xFF ] << 24 );

		X3 = *RK++ ^ \
		( (ms_u32) FSb[ ( Y3       ) & 0xFF ]       ) ^
		( (ms_u32) FSb[ ( Y0 >>  8 ) & 0xFF ] <<  8 ) ^
		( (ms_u32) FSb[ ( Y1 >> 16 ) & 0xFF ] << 16 ) ^
		( (ms_u32) FSb[ ( Y2 >> 24 ) & 0xFF ] << 24 );
	}

	PUT_ULONG_LE( X0, output,  0 );
	PUT_ULONG_LE( X1, output,  4 );
	PUT_ULONG_LE( X2, output,  8 );
	PUT_ULONG_LE( X3, output, 12 );
}

/*
 * AES-CBC buffer encryption/decryption
 */
static void msaes_innerapi_cryptCBC( ST_CA_AES_CONTEXT *ctx, ms_s32 mode,  ms_s32 length,
                    							ms_u08 iv[16], ms_u08 *input, ms_u08 *output )
{
	ms_s32 i,tmp =0;
    
	ms_u08 temp[16];
	if((input==ms_null)||(output==ms_null)){
        return;
	}
	if( mode == CRYPT_TYPE_DECRYPT ){
		while( length > 0 ){
			memcpy( temp, input, 16 );
			msaes_innerapi_cryptECB( ctx, mode, input, output );

			for( i = 0; i < 16; i++ ){
				output[i] = (ms_u08)( output[i] ^ iv[i] );
			}

			memcpy( iv, temp, 16 );
			if(tmp>200){
				tmp=0;
				break;
			}
			input  += 16;
			output += 16;
			length -= 16;
			tmp++;
		}
	}else{
		while( length > 0 ){
			for( i = 0; i < 16; i++ ){
				output[i] = (ms_u08)( input[i] ^ iv[i] );
			}

			msaes_innerapi_cryptECB( ctx, mode, output, output );
			memcpy( iv, output, 16 );
			if(tmp>200){
				tmp=0;
				break;
			}
			input  += 16;
			output += 16;
			length -= 16;
			tmp++;
		}
	}
}

/*
 * AES-CFB128 buffer encryption/decryption
 */
static void msaes_innerapi_cryptCFB128( ST_CA_AES_CONTEXT *ctx, ms_s32 mode, ms_s32 length, ms_s32 *iv_off,
                       							ms_u08 iv[16], ms_u08 *input,  ms_u08 *output )
{
	ms_s32 c, n = *iv_off;

	if( mode == CRYPT_TYPE_DECRYPT ){
		while( length-- ){
			if( n == 0 ){
				msaes_innerapi_cryptECB( ctx, CRYPT_TYPE_DECRYPT, iv, iv );
			}

			c = *input++;
			*output++ = (ms_u08)( c ^ iv[n] );
            
            if(n>16){
               printf("%s %d n:%d too big\n",__FUNCTION__,__LINE__,n);
               return;
            }
			iv[n] = (ms_u08) c;

			n = (n + 1) & 0x0F;
		}
	}else{
		while( length-- ){
			if( n == 0 ){
				msaes_innerapi_cryptECB( ctx, CRYPT_TYPE_DECRYPT, iv, iv );
			}

            if(n>16){
               printf("%s %d n:%d too big\n",__FUNCTION__,__LINE__,n);
               return;
            }
			iv[n] = *output++ = (ms_u08)( iv[n] ^ *input++ );
			n = (n + 1) & 0x0F;
		}
	}
	*iv_off = n;
}


ms_bool msaes_api_done(EN_CRYPT_TYPE entype, EN_CRYPT_MODE enmode, 
	ms_s08 * bufin, ms_s08 * bufout, ms_u32 datalen,  ms_s08* key,  ms_u08 keylen)   
{
	ms_s32 offset;
	ST_CA_AES_CONTEXT ctx;
	ms_u08 iv[16];
	ms_u08 ucKeyLenBits = 0;
	ms_s32 i = 0, j = 0;
	ms_bufcheckRetDes(ms_false, bufin, "bufin")
	ms_bufcheckRetDes(ms_false, bufout, "bufout")
	ms_bufcheckRetDes(ms_false, key, "key")

	if ((keylen != 32)&& (keylen != 16) && (keylen != 24)){
		ms_errRet(ms_false, "Keylen(%d) must be 32,16 or 24", keylen);
	}
	if ((datalen%16) != 0){
		ms_errRet(ms_false, "datalen(%d) must be 16*N", datalen);
	}
	ucKeyLenBits = keylen<<3;
	if (enmode == CRYPT_MODE_ECB){
		if (entype == CRYPT_TYPE_ENCRYPT){
			msaes_innerapi_setEncKey( &ctx, key, ucKeyLenBits );
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16)  
			{
				msaes_innerapi_cryptECB( &ctx, entype, bufin, bufout );
			}
		}else if (entype == CRYPT_TYPE_DECRYPT){
			msaes_innerapi_setDecKey( &ctx, key, ucKeyLenBits);
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16) {
				msaes_innerapi_cryptECB( &ctx, entype, bufin, bufout );
			}
		}
	}else if (enmode == CRYPT_MODE_CBC){
		if (entype == CRYPT_TYPE_ENCRYPT){
			msaes_innerapi_setEncKey( &ctx, key, ucKeyLenBits);
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16) {
				msaes_innerapi_cryptCBC( &ctx, entype, 16, iv, bufin, bufout );
			}
		}else if (entype == CRYPT_TYPE_DECRYPT){
			msaes_innerapi_setDecKey( &ctx, key, ucKeyLenBits );
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16) {
				msaes_innerapi_cryptCBC( &ctx, entype, 16, iv, bufin, bufout );
			}
		}
	}else{
		msaes_innerapi_setEncKey( &ctx, key, ucKeyLenBits );
		if (entype == CRYPT_TYPE_ENCRYPT){	
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16)  
			{
			 	msaes_innerapi_cryptCFB128( &ctx, entype, 64, &offset, iv, bufin, bufout );
			}
		}else if (entype == CRYPT_TYPE_DECRYPT){
			for(i=0,j=datalen>>4;i<j;++i,bufout+=16,bufin+=16) {
			 	msaes_innerapi_cryptCFB128( &ctx, entype, 64, &offset, iv, bufin, bufout );
			}
		}
	}
	return ms_true;
}
#undef MSAES_C
